/* ---------------------------------------------------------------------
 *
 * -- Privacy-preserving Sharing of Sensitive information Toolkit (PSST)
 *    (C) Copyright 2010 All Rights Reserved
 *
 * -- PSST routines -- Version 1.0.1 -- April, 2010
 * FNP corresponds to protocol FNP04 in the paper
 *
 * Author         : Yanbin Lu
 * Security and Privacy Research Outfit (SPROUT),
 * University of California, Irvine
 *
 * -- Bug report: send email to yanbinl@uci.edu
 * ---------------------------------------------------------------------
 *
 * -- Copyright notice and Licensing terms:
 *
 *  Redistribution  and  use in  source and binary forms, with or without
 *  modification, are  permitted provided  that the following  conditions
 *  are met:
 *
 * 1. Redistributions  of  source  code  must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce  the above copyright
 *    notice,  this list of conditions, and the  following disclaimer in
 *    the documentation and/or other materials provided with the distri-
 *    bution.
 * 3. The name of the University,  the SPROUT group,  or the names of its
 *    authors  may not be used to endorse or promote products deri-
 *    ved from this software without specific written permission.
 *
 * -- Disclaimer:
 *
 * THIS  SOFTWARE  IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES,  INCLUDING,  BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE UNIVERSITY
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,  INDIRECT, INCIDENTAL, SPE-
 * CIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
 * TO,  PROCUREMENT  OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
 * OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEO-
 * RY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT  (IN-
 * CLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

#include "fnp_psi_client.h"
#include "memreuse.h"
#include "mpz_tool.h"
#include "crypto_tool.h"
#include "basic_tool.h"
#include <stdio.h>
#include <math.h>

FNP_PSI_Client::FNP_PSI_Client()
{
  m_iNumIterations = 1;
  m_g = Memreuse::New();
  m_h = Memreuse::New();
  m_p = Memreuse::New();
  m_q = Memreuse::New();
  m_x = Memreuse::New();
  m_k = Memreuse::New();
}

void FNP_PSI_Client::LoadData()
{
  m_set.clear();
  m_rawset.clear();

  ifstream infile(m_strInputFile.c_str());
  if (!infile)
	{
	  cerr<<"open "<<m_strInputFile<<" error"<<endl;
	  return;
	}
  string element;
  while (getline(infile, element) && infile.good())
	{
	  m_rawset.push_back(element);
	  mpz_t * e = Memreuse::New();
	  CryptoTool::FDH(e, element, m_p);
	  m_set.push_back(e);
	}
}

void FNP_PSI_Client::Setup(int nbits)
{
  stringstream keyfile;
  keyfile<<"dsa."<<nbits<<".priv";
  BasicTool::Load_ElGamal_Priv_Key(m_g, m_p, m_q, m_x, m_h, keyfile.str());
  LoadData();

  rn.Initialize(1024);
}

// choose any n-power elements and sum their products
// sign: if power is even, the sign is positive
void FNP_PSI_Client::Coefficient(mpz_t * result, size_t original_power, size_t power, size_t j, mpz_t * product)
{
  if (power == m_set.size())
	{
      if ((original_power) % 2 == 0)
        {
          mpz_add(*result, *result, *product);
        }
      else
        {
          mpz_sub(*result, *result, *product);
        }
	  mpz_mod(*result, *result, *m_q);
	  return;
	}
  if (power < j)
	return;
  mpz_t * tmp = Memreuse::New();
  mpz_set(*tmp, *product);
  for (size_t i = j; i < m_set.size(); i++)
	{
	  mpz_mul(*product, *product, *m_set[i]);
	  mpz_mod(*product, *product, *m_q);
	  Coefficient(result, original_power, power+1, i+1, product);
	  mpz_set(*product, *tmp);
	}
  Memreuse::Delete(tmp);
}
// vecRoot holds negative m_set[i]
void FNP_PSI_Client::ComputeCoefficient(vector<mpz_t *>& vecCoefficient, const vector<mpz_t *>& vecRoot)
{
  vecCoefficient.resize(vecRoot.size() + 1);
  for (size_t i = 0; i < vecCoefficient.size(); i++)
    {
	  if (vecCoefficient[i] == NULL)
		vecCoefficient[i] = Memreuse::New();
      mpz_set_ui(*vecCoefficient[i], 0);
    }

  mpz_set_ui(*vecCoefficient[1], 1);
  mpz_set(*vecCoefficient[0], *vecRoot[0]);
  //gmp_printf("%Zx\n", vecCoefficient[0]);
  for (size_t i = 1; i < vecRoot.size(); i++)
    {
      for (size_t j = i + 1; j > 0; j--)
        {
          // vecCoefficnet[j]=vecCoefficnet[j-1] + root[i]*vecCoefficnet[j]
          mpz_mul(*vecCoefficient[j], *vecCoefficient[j], *vecRoot[i]);
          mpz_add(*vecCoefficient[j], *vecCoefficient[j-1], *vecCoefficient[j]);
          mpz_mod(*vecCoefficient[j], *vecCoefficient[j], *m_q);
        }
      mpz_mul(*vecCoefficient[0], *vecCoefficient[0], *vecRoot[i]);
      mpz_mod(*vecCoefficient[0], *vecCoefficient[0], *m_q);
    }
}

size_t FNP_PSI_Client::WhichBin(mpz_t * id, mpz_t * k)
{
  mpz_t * newid = Memreuse::New();
  mpz_xor (*newid, *id, *k);
  mpz_mod_ui (*newid, *newid, m_B);
  unsigned long key = mpz_get_ui (*newid);
  Memreuse::Delete(newid);
  return key;
}

int FNP_PSI_Client::Request1(vector<string>& output)
{
  output.clear();
  if (m_set.size() > 1)
	m_B = (size_t) ceil(m_set.size() / log(log((double)m_set.size())));
  else
	m_B = 1;
  m_vecPoly.resize(m_B);
  rn.Get_Z_Range(m_k, m_p);		// changing mapping each time


  for (size_t i = 0; i < m_set.size(); i++)
	{
	  mpz_t * invm = Memreuse::New();
	  size_t b = WhichBin(m_set[i], m_k);
	  // negative for computing coefficient
      mpz_neg(*invm, *m_set[i]);
      mpz_mod(*invm, *invm, *m_q);

	  m_vecPoly[b].push_back(invm);
	}

  m_M = 0;						// maximum number of elements in one bin
  for (size_t i = 0; i < m_vecPoly.size(); i++)
	{
	  if (m_vecPoly[i].size() > m_M)
		{
		  m_M = m_vecPoly[i].size();
		}
	}
  // pad zero
  mpz_t * tmp = Memreuse::New();
  mpz_set_ui(*tmp, 0);
  for (size_t i = 0; i < m_vecPoly.size(); i++)
	{
	  // make each bin have M roots
	  while (m_vecPoly[i].size() < m_M)
		{
		  m_vecPoly[i].push_back(tmp);
		}
	}
  stringstream ss;
  ss<<m_M;
  output.push_back(ss.str());
  ss.str("");
  ss<<m_B;
  output.push_back(ss.str());
  output.push_back(MpzTool::String_from_mpz(NULL, m_k));
  mpz_t * elgamal_l = Memreuse::New();
  mpz_t * elgamal_r = Memreuse::New();
  for (size_t i = 0; i < m_vecPoly.size(); i++)
	{
	  // m_coefficient will be initialized inside ComputeCoefficient
	  ComputeCoefficient(m_coefficient, m_vecPoly[i]);
	  for (size_t i = 0; i < m_coefficient.size(); i++)
		{
		  //gmp_printf("%Zx\n", *m_coefficient[i]);
		  // encrypt the coefficient
		  mpz_powm(*m_coefficient[i], *m_g, *m_coefficient[i], *m_p);		// g^m
		  CryptoTool::ElGamalEnc(elgamal_l, elgamal_r, m_coefficient[i], m_g, m_p, m_q, m_h);
		  output.push_back(MpzTool::String_from_mpz(NULL, elgamal_l));
		  output.push_back(MpzTool::String_from_mpz(NULL, elgamal_r));
		}
	}
  Memreuse::Delete(elgamal_l);
  Memreuse::Delete(elgamal_r);	  
  mpz_t * gtom = Memreuse::New();
  for (size_t i = 0; i < m_set.size(); i++)
	{
	  mpz_powm(*gtom, *m_g, *m_set[i], *m_p);
	  m_mapTp[MpzTool::Mpz_to_string(gtom)] = i;
	}
  Memreuse::Delete(gtom);
  return 0;
}

int FNP_PSI_Client::OnResponse1(const vector<string>& input)
{
  mpz_t * el = Memreuse::New();
  mpz_t * er = Memreuse::New();
  mpz_t * result = Memreuse::New();
  m_vecResult.clear();
  map<string, size_t>::iterator it;
  for (size_t i = 0; i < input.size(); i = i+2)
	{
	  MpzTool::Mpz_from_string(el, input[i]);
	  MpzTool::Mpz_from_string(er, input[i+1]);
	  // decrypt
	  CryptoTool::ElGamalDec(result, el, er, m_g, m_p, m_q, m_x);
	  if ((it = m_mapTp.find(MpzTool::Mpz_to_string(result))) != m_mapTp.end())
        {
          m_vecResult.push_back(it->second);
        }
	}

  Memreuse::Delete(el);
  Memreuse::Delete(er);
  Memreuse::Delete(result);
  return 0;
}
